import argparse
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from tqdm import tqdm

from dataset import ImageFileDataset
from vqvae import VQVAE

import converter

# TODO:
#  1. train a auto-encoder with binary constraint: Loss(binary)=-2*|sigmoid(z)-0.5|
#  2. dataset generated from seeprettyface model
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--size', type=int, default=256)
    parser.add_argument('--ckpt', type=str, default='./vqvae_560.pt')
    parser.add_argument('--name', type=str, default='test')
    parser.add_argument('--source', type=str, default='../Datasets/Text/')
    parser.add_argument('--target', type=str, default='../Datasets/SR/Text/')
    parser.add_argument('--cache', type=str, default='../Datasets/Cache')
    args = parser.parse_args()
    device = 'cpu'
    transform = transforms.Compose(
        [
            transforms.Resize(args.size),
            transforms.CenterCrop(args.size),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        ]
    )
    dataset = ImageFileDataset(args.source, transform=transform)
    loader = DataLoader(dataset, batch_size=1, shuffle=False)

    model = VQVAE()
    model.load_state_dict(torch.load(args.ckpt, map_location='cpu'))
    model = model.to(device)
    model.eval()

    index = 0
    pbar = tqdm(loader)
    for img, _, filename in pbar:
        img = img.to(device)
        _, _, _, id_t, id_b = model.encode(img)
        id_t = id_t.detach()
        id_b = id_b.detach()
        res = model.decode_code(id_t, id_b).clamp(-1, 1)
        im_ = converter.tensor2im(res.squeeze(0).detach().cpu())
        im_ = Image.fromarray(im_)
        im_.save(args.cache + '/%dx%d.png' % (im_.width, im_.height))
        exit(0)
        index += 1
        input()